import numpy as np
import pandas as pd
from tensorflow import keras
from tensorflow.keras import layers
from sklearn.preprocessing import MinMaxScaler
import logging

logger = logging.getLogger(__name__)


class LSTMModel:
    """LSTM深度学习时间序列预测模型"""

    def __init__(self, lookback=168):  # 默认使用过去7天数据
        self.model = None
        self.scaler = MinMaxScaler(feature_range=(0, 1))
        self.lookback = lookback  # 时间窗口大小

    def build_model(self, input_shape):
        """构建LSTM模型"""
        model = keras.Sequential([
            layers.LSTM(128, return_sequences=True, input_shape=input_shape),
            layers.Dropout(0.2),
            layers.LSTM(64, return_sequences=True),
            layers.Dropout(0.2),
            layers.LSTM(32),
            layers.Dropout(0.2),
            layers.Dense(24, activation='relu'),  # 预测未来24小时
            layers.Dense(24)
        ])

        model.compile(
            optimizer=keras.optimizers.Adam(learning_rate=0.001),
            loss='mse',
            metrics=['mae']
        )

        return model

    def prepare_sequences(self, data):
        """准备LSTM训练序列"""
        scaled_data = self.scaler.fit_transform(data.reshape(-1, 1))

        X, y = [], []
        for i in range(self.lookback, len(scaled_data) - 24):
            X.append(scaled_data[i - self.lookback:i, 0])
            y.append(scaled_data[i:i + 24, 0])

        return np.array(X), np.array(y)

    def predict(self, df, forecast_hours=24, confidence_level=0.95):
        """
        使用LSTM进行预测

        Args:
            df: DataFrame with columns ['timestamp', 'value']
            forecast_hours: 预测未来多少小时
            confidence_level: 置信区间水平

        Returns:
            预测结果列表
        """
        try:
            logger.info("开始LSTM预测...")

            # 准备数据
            values = df['value'].values

            if len(values) < self.lookback + 24:
                raise ValueError(f"数据量不足，至少需要{self.lookback + 24}个数据点")

            # 准备训练数据
            X, y = self.prepare_sequences(values)
            X = X.reshape((X.shape[0], X.shape[1], 1))

            # 构建并训练模型
            self.model = self.build_model((X.shape[1], 1))

            # 训练模型（使用较少的epoch以加快速度，实际应用可以增加）
            history = self.model.fit(
                X, y,
                epochs=50,
                batch_size=32,
                validation_split=0.1,
                verbose=0,
                callbacks=[
                    keras.callbacks.EarlyStopping(
                        monitor='val_loss',
                        patience=5,
                        restore_best_weights=True
                    )
                ]
            )

            # 准备预测输入（使用最近的数据）
            last_sequence = self.scaler.transform(
                values[-self.lookback:].reshape(-1, 1)
            )
            last_sequence = last_sequence.reshape((1, self.lookback, 1))

            # 执行预测
            predictions_scaled = self.model.predict(last_sequence, verbose=0)
            predictions = self.scaler.inverse_transform(predictions_scaled)

            # 计算置信区间（使用预测误差的标准差）
            train_predictions = self.model.predict(X, verbose=0)
            train_predictions = self.scaler.inverse_transform(train_predictions)
            actual_values = self.scaler.inverse_transform(y)

            errors = actual_values - train_predictions
            std_error = np.std(errors)
            z_score = 1.96 if confidence_level == 0.95 else 2.576  # 95% or 99%

            # 生成时间戳
            last_timestamp = pd.to_datetime(df['timestamp'].iloc[-1])
            future_timestamps = pd.date_range(
                start=last_timestamp + pd.Timedelta(hours=1),
                periods=min(forecast_hours, 24),
                freq='H'
            )

            # 构建返回结果
            results = []
            for i, ts in enumerate(future_timestamps):
                pred_value = max(0, int(predictions[0][i]))
                results.append({
                    'timestamp': ts.strftime('%Y-%m-%d %H:%M:%S'),
                    'predicted_value': pred_value,
                    'confidence_lower': max(0, int(pred_value - z_score * std_error)),
                    'confidence_upper': int(pred_value + z_score * std_error)
                })

            logger.info(f"LSTM预测完成，生成{len(results)}个预测点")
            return results

        except Exception as e:
            logger.error(f"LSTM预测失败: {str(e)}")
            raise